! (C) Copyright 2008- ECMWF.
! (C) Copyright 2008- Meteo-France.
! 
! This software is licensed under the terms of the Apache Licence Version 2.0
! which can be obtained at http://www.apache.org/licenses/LICENSE-2.0.
! In applying this licence, ECMWF does not waive the privileges and immunities
! granted to it by virtue of its status as an intergovernmental organisation
! nor does it submit to any jurisdiction.
!

MODULE GPNORM_TRANS_CTL_MOD
CONTAINS
SUBROUTINE GPNORM_TRANS_CTL(PGP,KFIELDS,KPROMA,PAVE,PMIN,PMAX,LDAVE_ONLY,PW)


!**** *GPNORM_TRANS_CTL* - calculate grid-point norms

!     Purpose.
!     --------
!        calculate grid-point norms using a 2 stage (NPRTRV,NPRTRW) communication rather
!        than an approach using a more expensive global gather collective communication

!**   Interface.
!     ----------
!     CALL GPNORM_TRANS_CTL(...)

!     Explicit arguments :
!     --------------------
!     PGP(:,:,:) - gridpoint fields (input)
!                  PGP is  dimensioned (NPROMA,KFIELDS,NGPBLKS) where
!                  NPROMA is the blocking factor, KFIELDS the total number
!                  of fields and NGPBLKS the number of NPROMA blocks.
!     KFIELDS     - number of fields (input)
!                   (these do not have to be just levels)
!     KPROMA      - required blocking factor (input)
!     PAVE        - average (output)
!     PMIN        - minimum (input/output)
!     PMAX        - maximum (input/output)
!     LDAVE_ONLY  - T : PMIN and PMAX already contain local MIN and MAX
!

!     Author.
!     -------
!        George Mozdzynski *ECMWF*

!     Modifications.
!     --------------
!        Original : 19th Sept 2008
!        R. El Khatib 07-08-2009 Optimisation directive for NEC
!        R. El Khatib 16-Sep-2019 merge with LAM code
!        R. El Khatib 02-Jun-2022 Optimization/Cleaning
!     ------------------------------------------------------------------

USE PARKIND1  ,ONLY : JPIM     ,JPRB, JPRD

!ifndef INTERFACE

USE TPM_GEN         ,ONLY : NOUT
USE TPM_DIM         ,ONLY : R
USE TPM_TRANS       ,ONLY : LGPNORM, NGPBLKS, NPROMA
USE TPM_DISTR       ,ONLY : D, NPRCIDS, NPRTRV, NPRTRW, MYSETV, MYSETW
USE TPM_GEOMETRY    ,ONLY : G
USE SET_RESOL_MOD   ,ONLY : SET_RESOL
USE TRGTOL_MOD      ,ONLY : TRGTOL
USE SET2PE_MOD      ,ONLY : SET2PE
USE MPL_MODULE      ,ONLY : MPL_RECV, MPL_SEND, JP_BLOCKING_STANDARD
USE ABORT_TRANS_MOD ,ONLY : ABORT_TRANS
USE YOMHOOK         ,ONLY : LHOOK,   DR_HOOK,  JPHOOK

!endif INTERFACE

IMPLICIT NONE

! Declaration of arguments

REAL(KIND=JPRB)   ,INTENT(IN)    :: PGP(:,:,:)
REAL(KIND=JPRB)   ,INTENT(OUT)   :: PAVE(:)
REAL(KIND=JPRB)   ,INTENT(INOUT) :: PMIN(:)
REAL(KIND=JPRB)   ,INTENT(INOUT) :: PMAX(:)
INTEGER(KIND=JPIM),INTENT(IN)    :: KFIELDS
INTEGER(KIND=JPIM),INTENT(IN)    :: KPROMA
LOGICAL           ,INTENT(IN)    :: LDAVE_ONLY
REAL(KIND=JPRD)   ,INTENT(IN)    :: PW(R%NDGL)

!ifndef INTERFACE

! Local variables
REAL(KIND=JPHOOK)  :: ZHOOK_HANDLE
INTEGER(KIND=JPIM) :: IUBOUND(4)
INTEGER(KIND=JPIM) :: IVSET(KFIELDS)
INTEGER(KIND=JPIM),ALLOCATABLE :: IVSETS(:)
INTEGER(KIND=JPIM),ALLOCATABLE :: IVSETG(:,:)
REAL(KIND=JPRB),ALLOCATABLE :: ZGTF(:,:)
REAL(KIND=JPRD),ALLOCATABLE :: ZAVE(:,:)
REAL(KIND=JPRB),ALLOCATABLE :: ZMINGL(:,:)
REAL(KIND=JPRB),ALLOCATABLE :: ZMAXGL(:,:)
REAL(KIND=JPRB),ALLOCATABLE :: ZMIN(:)
REAL(KIND=JPRB),ALLOCATABLE :: ZMAX(:)
REAL(KIND=JPRD),ALLOCATABLE :: ZAVEG(:,:)
REAL(KIND=JPRB),ALLOCATABLE :: ZMING(:)
REAL(KIND=JPRB),ALLOCATABLE :: ZMAXG(:)
REAL(KIND=JPRD),ALLOCATABLE :: ZSND(:)
REAL(KIND=JPRD),ALLOCATABLE :: ZRCV(:)
INTEGER(KIND=JPIM) :: J,JGL,IGL,JL,JF,IF_GP,IF_SCALARS_G,IF_FS,JSETV,JSETW,IWLATS
INTEGER(KIND=JPIM) :: IPROC,ITAG,ILEN,ILENR,IBEG,IEND,IND

!     ------------------------------------------------------------------
IF (LHOOK) CALL DR_HOOK('GPNORM_TRANS_CTL',0,ZHOOK_HANDLE)

! Set defaults

NPROMA = KPROMA
NGPBLKS = (D%NGPTOT-1)/NPROMA+1

! Consistency checks

IUBOUND(1:3)=UBOUND(PGP)
IF(IUBOUND(1) < NPROMA) THEN
  WRITE(NOUT,*)'GPNORM_TRANS_CTL:FIRST DIM. OF PGP TOO SMALL ',IUBOUND(1),NPROMA
  CALL ABORT_TRANS('GPNORM_TRANS_CTL:FIRST DIMENSION OF PGP TOO SMALL ')
ENDIF
IF(IUBOUND(2) < KFIELDS) THEN
  WRITE(NOUT,*)'GPNORM_TRANS_CTL:SEC. DIM. OF PGP TOO SMALL ',IUBOUND(2),KFIELDS
  CALL ABORT_TRANS('GPNORM_TRANS_CTL:SECOND DIMENSION OF PGP TOO SMALL ')
ENDIF
IF(IUBOUND(3) < NGPBLKS) THEN
  WRITE(NOUT,*)'GPNORM_TRANS_CTL:THIRD DIM. OF PGP TOO SMALL ',IUBOUND(3),NGPBLKS
  CALL ABORT_TRANS('GPNORM_TRANS_CTL:THIRD DIMENSION OF PGP TOO SMALL ')
ENDIF


IF_GP=KFIELDS
IF_SCALARS_G=0

IF_FS=0
DO J=1,KFIELDS
  IVSET(J)=MOD(J-1,NPRTRV)+1
  IF(IVSET(J)==MYSETV)THEN
    IF_FS=IF_FS+1
  ENDIF
ENDDO

ALLOCATE(IVSETS(NPRTRV))
IVSETS(:)=0
DO J=1,KFIELDS
  IVSETS(IVSET(J))=IVSETS(IVSET(J))+1
ENDDO
ALLOCATE(IVSETG(NPRTRV,MAXVAL(IVSETS(:))))
IVSETG(:,:)=0
IVSETS(:)=0
DO J=1,KFIELDS
  IVSETS(IVSET(J))=IVSETS(IVSET(J))+1
  IVSETG(IVSET(J),IVSETS(IVSET(J)))=J
ENDDO

ALLOCATE(ZGTF(IF_FS,D%NLENGTF))
IF (SIZE(ZGTF) > 0) ZGTF(1,1)=0._JPRB ! force allocation right here, not inside an omp region below
LGPNORM=.TRUE.
CALL TRGTOL(ZGTF,IF_FS,IF_GP,IF_SCALARS_G,IVSET,PGP=PGP)
LGPNORM=.FALSE.

IBEG=1
IEND=D%NDGL_FS

ALLOCATE(ZAVE(IF_FS,IBEG:IEND))
ALLOCATE(ZMIN(IF_FS))
ALLOCATE(ZMAX(IF_FS))
IF(.NOT.LDAVE_ONLY)THEN
  ALLOCATE(ZMINGL(IF_FS,IBEG:IEND))
  ALLOCATE(ZMAXGL(IF_FS,IBEG:IEND))
ENDIF

IF( IF_FS > 0 )THEN

  ZAVE(:,:)=0.0_JPRB

  IF(.NOT.LDAVE_ONLY)THEN
    DO JF=1,IF_FS
      ZMINGL(JF,:)=ZGTF(JF,D%NSTAGTF(1)+1)
      ZMAXGL(JF,:)=ZGTF(JF,D%NSTAGTF(1)+1)
    ENDDO
  ENDIF

! FIRST DO SUMS IN EACH FULL LATITUDE

CALL GSTATS(1429,0)
!$OMP PARALLEL DO SCHEDULE(STATIC) PRIVATE(JGL,IGL,JF,JL)
  DO JGL=IBEG,IEND
    IGL = D%NPTRLS(MYSETW) + JGL - 1
!CDIR NOLOOPCHG
    DO JF=1,IF_FS
!DIR$ NEXTSCALAR
      DO JL=1,G%NLOEN(IGL)
        ZAVE(JF,JGL)=ZAVE(JF,JGL)+ZGTF(JF,D%NSTAGTF(JGL)+JL)
      ENDDO
      IF(.NOT.LDAVE_ONLY)THEN
        DO JL=1,G%NLOEN(IGL)
          ZMINGL(JF,JGL)=MIN(ZMINGL(JF,JGL),ZGTF(JF,D%NSTAGTF(JGL)+JL))
          ZMAXGL(JF,JGL)=MAX(ZMAXGL(JF,JGL),ZGTF(JF,D%NSTAGTF(JGL)+JL))
        ENDDO
      ENDIF
    ENDDO
  ENDDO
!$OMP END PARALLEL DO
CALL GSTATS(1429,1)

  IF(.NOT.LDAVE_ONLY)THEN
    DO JF=1,IF_FS
      ZMIN(JF)=MINVAL(ZMINGL(JF,:))
      ZMAX(JF)=MAXVAL(ZMAXGL(JF,:))
    ENDDO
    DEALLOCATE(ZMINGL)
    DEALLOCATE(ZMAXGL)
  ENDIF

  DO JGL=IBEG,IEND
    IGL = D%NPTRLS(MYSETW) + JGL - 1
    DO JF=1,IF_FS
      ZAVE(JF,JGL)=ZAVE(JF,JGL)*REAL(PW(IGL),JPRB)/G%NLOEN(IGL)
    ENDDO
  ENDDO

ENDIF

! IT IS IMPORTANT THAT SUMS ARE NOW DONE IN LATITUDE ORDER
ALLOCATE(ZAVEG(R%NDGL,KFIELDS))
ALLOCATE(ZMING(KFIELDS))
ALLOCATE(ZMAXG(KFIELDS))

ZAVEG(:,:)=0.0_JPRB

DO JF=1,IF_FS
  DO JGL=IBEG,IEND
    IGL = D%NPTRLS(MYSETW) + JGL - 1
    ZAVEG(IGL,IVSETG(MYSETV,JF))=ZAVEG(IGL,IVSETG(MYSETV,JF))+ZAVE(JF,JGL)
  ENDDO
ENDDO

IF(LDAVE_ONLY)THEN
  ZMING(:)=PMIN(:)
  ZMAXG(:)=PMAX(:)
ELSE
  DO JF=1,IF_FS
    ZMING(IVSETG(MYSETV,JF))=ZMIN(JF)
    ZMAXG(IVSETG(MYSETV,JF))=ZMAX(JF)
  ENDDO
ENDIF

! RECEIVE ABOVE FROM OTHER NPRTRV SETS FOR SAME LATS BUT DIFFERENT FIELDS
ITAG=123

CALL GSTATS(815,0)

IF( MYSETV==1 )THEN

  DO JSETV=2,NPRTRV
    IF(LDAVE_ONLY)THEN
      ILEN=D%NDGL_FS*IVSETS(JSETV)+2*KFIELDS
    ELSE
      ILEN=(D%NDGL_FS+2)*IVSETS(JSETV)
    ENDIF
    IF(ILEN > 0)THEN
      ALLOCATE(ZRCV(ILEN))
      CALL SET2PE(IPROC,0,0,MYSETW,JSETV)
      CALL MPL_RECV(ZRCV(:),KSOURCE=NPRCIDS(IPROC),KTAG=ITAG,&
        &KMP_TYPE=JP_BLOCKING_STANDARD,KOUNT=ILENR,CDSTRING='GPNORM_TRANS_CTL:V')
      IF(ILENR /= ILEN)THEN
        CALL ABOR1('GPNORM_TRANS_CTL:ILENR /= ILEN')
      ENDIF
      IND=0
      DO JF=1,IVSETS(JSETV)
        DO JGL=IBEG,IEND
          IGL = D%NPTRLS(MYSETW) + JGL - 1
          IND=IND+1
          ZAVEG(IGL,IVSETG(JSETV,JF))=ZRCV(IND)
        ENDDO
        IF(.NOT.LDAVE_ONLY)THEN
          IND=IND+1
          ZMING(IVSETG(JSETV,JF))=ZRCV(IND)
          IND=IND+1
          ZMAXG(IVSETG(JSETV,JF))=ZRCV(IND)
        ENDIF
      ENDDO
      IF(LDAVE_ONLY)THEN
        DO JF=1,KFIELDS
          IND=IND+1
          ZMING(JF)=MIN(ZMING(JF),REAL(ZRCV(IND),KIND=JPRB))
          IND=IND+1
          ZMAXG(JF)=MAX(ZMAXG(JF),REAL(ZRCV(IND),KIND=JPRB))
        ENDDO
      ENDIF
      DEALLOCATE(ZRCV)
    ENDIF
  ENDDO

ELSE

  IF(LDAVE_ONLY)THEN
    ILEN=D%NDGL_FS*IVSETS(MYSETV)+2*KFIELDS
  ELSE
    ILEN=(D%NDGL_FS+2)*IVSETS(MYSETV)
  ENDIF
  IF(ILEN > 0)THEN
    CALL SET2PE(IPROC,0,0,MYSETW,1)
    ALLOCATE(ZSND(ILEN))
    IND=0
    DO JF=1,IF_FS
      DO JGL=IBEG,IEND
        IGL = D%NPTRLS(MYSETW) + JGL - 1
        IND=IND+1
        ZSND(IND)=ZAVEG(IGL,IVSETG(MYSETV,JF))
       ENDDO
      IF(.NOT.LDAVE_ONLY)THEN
        IND=IND+1
        ZSND(IND)=ZMING(IVSETG(MYSETV,JF))
        IND=IND+1
        ZSND(IND)=ZMAXG(IVSETG(MYSETV,JF))
      ENDIF
    ENDDO
    IF(LDAVE_ONLY)THEN
      DO JF=1,KFIELDS
        IND=IND+1
        ZSND(IND)=PMIN(JF)
        IND=IND+1
        ZSND(IND)=PMAX(JF)
      ENDDO
    ENDIF
    CALL MPL_SEND(ZSND(:),KDEST=NPRCIDS(IPROC),KTAG=ITAG,&
      &KMP_TYPE=JP_BLOCKING_STANDARD,CDSTRING='GPNORM_TRANS_CTL:V')
    DEALLOCATE(ZSND)
  ENDIF

ENDIF

! FINALLY RECEIVE CONTRIBUTIONS FROM OTHER NPRTRW SETS

IF( MYSETV == 1 )THEN

  IF( MYSETW == 1 )THEN

    DO JSETW=2,NPRTRW
      IWLATS=D%NULTPP(JSETW)
      IBEG=1
      IEND=IWLATS
      IF(LDAVE_ONLY)THEN
        ILEN=IWLATS*KFIELDS+2*KFIELDS
      ELSE
        ILEN=(IWLATS+2)*KFIELDS
      ENDIF
      IF(ILEN > 0 )THEN
        ALLOCATE(ZRCV(ILEN))
        CALL SET2PE(IPROC,0,0,JSETW,1)
        CALL MPL_RECV(ZRCV(:),KSOURCE=NPRCIDS(IPROC),KTAG=ITAG,&
          &KMP_TYPE=JP_BLOCKING_STANDARD,KOUNT=ILENR,CDSTRING='GPNORM_TRANS_CTL:W')
        IF(ILENR /= ILEN)THEN
          CALL ABOR1('GPNORM_TRANS_CTL:ILENR /= ILEN')
        ENDIF
        IND=0
        DO JF=1,KFIELDS
          DO JGL=IBEG,IEND
            IGL = D%NPTRLS(JSETW) + JGL - 1
            IND=IND+1
            ZAVEG(IGL,JF)=ZRCV(IND)
          ENDDO
          IF(.NOT.LDAVE_ONLY)THEN
            IND=IND+1
            ZMING(JF)=MIN(ZMING(JF),REAL(ZRCV(IND),KIND=JPRB))
            IND=IND+1
            ZMAXG(JF)=MAX(ZMAXG(JF),REAL(ZRCV(IND),KIND=JPRB))
          ENDIF
        ENDDO
        IF(LDAVE_ONLY)THEN
          DO JF=1,KFIELDS
            IND=IND+1
            ZMING(JF)=MIN(ZMING(JF),REAL(ZRCV(IND),KIND=JPRB))
            IND=IND+1
            ZMAXG(JF)=MAX(ZMAXG(JF),REAL(ZRCV(IND),KIND=JPRB))
          ENDDO
        ENDIF
        DEALLOCATE(ZRCV)
      ENDIF
    ENDDO

  ELSE

    IF(LDAVE_ONLY)THEN
      ILEN=D%NDGL_FS*KFIELDS+2*KFIELDS
    ELSE
      ILEN=(D%NDGL_FS+2)*KFIELDS
    ENDIF
    IF(ILEN > 0)THEN
      CALL SET2PE(IPROC,0,0,1,1)
      ALLOCATE(ZSND(ILEN))
      IND=0
      DO JF=1,KFIELDS
        DO JGL=IBEG,IEND
          IGL = D%NPTRLS(MYSETW) + JGL - 1
          IND=IND+1
          ZSND(IND)=ZAVEG(IGL,JF)
        ENDDO
        IF(.NOT.LDAVE_ONLY)THEN
          IND=IND+1
          ZSND(IND)=ZMING(JF)
          IND=IND+1
          ZSND(IND)=ZMAXG(JF)
        ENDIF
      ENDDO
      IF(LDAVE_ONLY)THEN
        DO JF=1,KFIELDS
          IND=IND+1
          ZSND(IND)=ZMING(JF)
          IND=IND+1
          ZSND(IND)=ZMAXG(JF)
        ENDDO
      ENDIF
      CALL MPL_SEND(ZSND(:),KDEST=NPRCIDS(IPROC),KTAG=ITAG,&
        &KMP_TYPE=JP_BLOCKING_STANDARD,CDSTRING='GPNORM_TRANS_CTL:V')
      DEALLOCATE(ZSND)
    ENDIF

  ENDIF

ENDIF

CALL GSTATS(815,1)

IF( MYSETW == 1 .AND. MYSETV == 1 )THEN

  PAVE(:)=0.0_JPRB
  DO JGL=1,R%NDGL
    PAVE(:)=PAVE(:)+ZAVEG(JGL,:)
  ENDDO

  PMIN(:)=ZMING(:)
  PMAX(:)=ZMAXG(:)

ENDIF

DEALLOCATE(ZGTF)
DEALLOCATE(ZAVE)
DEALLOCATE(ZMIN)
DEALLOCATE(ZMAX)
DEALLOCATE(ZAVEG)
DEALLOCATE(ZMING)
DEALLOCATE(ZMAXG)
DEALLOCATE(IVSETS)
DEALLOCATE(IVSETG)

IF (LHOOK) CALL DR_HOOK('GPNORM_TRANS_CTL',1,ZHOOK_HANDLE)

!     ------------------------------------------------------------------

!endif INTERFACE


END SUBROUTINE GPNORM_TRANS_CTL
END MODULE GPNORM_TRANS_CTL_MOD
